/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: hdcsImpl实现头文件
 * Author: liutianming
 * Create: 2020-01-05
 */

#ifndef HDCS_IMPL_H
#define HDCS_IMPL_H

#include <hccl/hccl_types.h>
#include "hccl/base.h"
#include "base.h"
#include "log.h"
#include "memory_alloc_ring.h"
#include "op_hdcs_remote_lookup_executor.h"
#include "op_hdcs_remote_update_executor.h"
#include "transport_shm_event_pub.h"
#include "hdds_param.h"
#include "adapter_hal.h"
#include "alg_remote_lookup.h"
#include "alg_remote_update.h"
#include "hashtable/universal_concurrent_map.h"

namespace hccl {

class HdcsImpl {
public:
    explicit HdcsImpl();
    ~HdcsImpl();

    HcclResult Init(HdcsCsInitPara para, std::vector<u32> psRankId_);
    HcclResult DeInit();
    HcclResult RegTransport(HdcsRegTransportPara para);
    bool IsInited();

    HcclResult RemoteLookUp(HdcsEmbeddingServiceParam *lookupInfo);
    HcclResult IrecvLookupResponse(HcclRequest request);
    HcclResult RemoteUpdate(HdcsEmbeddingServiceParam *updateInfo);
    HcclResult IrecvUpdateResponse(HcclRequest request);
    HcclResult RemoteUpdateKeyReduce(HdcsEmbeddingServiceParam *updateInfo);
    HcclResult RemoteUpdateReduceSum(HdcsEmbeddingServiceParam *updateInfo);
    HcclResult RemoteUpdateSendRequest(HdcsEmbeddingServiceParam *updateInfo);
    HcclResult RemoteUpdateResetUniqueHandle(HdcsEmbeddingServiceParam *updateInfo);
    HcclResult RemoteUpdateRecvResponse(HdcsEmbeddingServiceParam *updateInfo);
    HcclResult TestSome(HcclRequest* request);
    HcclResult CollLookupKeysDuplicates(HdcsEmbeddingServiceParam *lookupInfo);
    HcclResult CollLookupSendKeys(HdcsLookupPipelineParams *pipelineParams);
    HcclResult CollLookupRecvValues(HdcsLookupPipelineParams *pipelineParams);
    HcclResult CollLookupRecoverValue(HdcsLookupPipelineParams *pipelineParams);
    HcclResult CollLookupWaitSendKeyFinish(HdcsLookupPipelineParams *pipelineParams);
    HcclResult CollLookupResetUniquesHandle(HdcsLookupPipelineParams *pipelineParams);

    HcclResult ChkServiceCancel(HdcsEmbeddingServiceParam *inferenceParams);
private:
    bool initFlag_;
    HcclComm comm_;
    LocklessRingMemoryAllocate<HcclOperator<EmbeddingServiceParam>> embeddingRequestPool_;

    // key : tag + cube index
    UniversalConcurrentMap<s32, std::unordered_map<s32, std::unique_ptr<OpExecutor>>> lookupOpExecutorMap_;

    // key : tag + cube index
    UniversalConcurrentMap<s32, std::unordered_map<s32, HcclOperator<EmbeddingServiceParam>*>> lookupParamsMap_;
    HdcsOpRemoteLookUpExecutor opRemoteLookUpExecutor_;
    HdcsOpRemoteUpdateExecutor opRemoteUpdateExecutor_;

    std::vector<u32> psRankId_;
    AlgRemoteLookUp lookupAlg_;
    AlgHdcsRemoteUpdate updateAlg_;

    // key: tag
    UniversalConcurrentMap<s32, HcclOperator<EmbeddingServiceParam>*> updateParamsMap_;
    UniversalConcurrentMap<s32, std::unique_ptr<HdcsOpRemoteUpdateExecutor>> updateOpExecutorMap_;
};

}  // end namespace hccl
#endif  // HDCS_IMPL_BASE_H
